# -*- coding: utf-8 -*-
# FileName:     zwave.py
# Time:         23/5/15 015 下午 3:32
# Author:       Zhou Hang
# Description:  对wave的包装
import wave
import matplotlib.pyplot as plt
import numpy as np

from src.utils.common import float2sample


def get_dtype_by_samplewidth(samplewidth: int):
    """
    根据采样精度返回对应dtype
    :param samplewidth:
    :return:
    """
    sample_width_tbl = {
        1: np.int8,
        2: np.int16,
        4: np.int32
    }
    assert samplewidth in sample_width_tbl.keys()
    return sample_width_tbl[samplewidth]


class WriteWave:
    def __init__(self, filepath: str):
        """
        设置默认的录音参数, 除非调用set_attr, 否则就使用默认参数
        :param filepath:
        """
        self.nchannels = 1  # 通道数
        self.samplewidth = 2  # 采样精度 2byte/16bit
        self.samplerate = 44100  # 采样率
        try:
            self.filepath = filepath
            self.wav_file = wave.open(self.filepath, "wb")
        except FileNotFoundError as e:
            print("文件名不合法，退出程序")
            assert False
        self.set_attr(self.nchannels, self.samplewidth, self.samplerate)

    def set_attr(self, nchannels: int, samplewidth: int, samplerate: int):
        """
        设置音频文件参数，以备后续写入，该方法已经在init中设置过默认参数了
        :param nchannels: 通道数
        :param samplewidth: 采样精度
        :param samplerate: 采样率
        :return:
        """
        self.nchannels = nchannels
        self.samplewidth = samplewidth
        self.samplerate = samplerate
        self.wav_file.setnchannels(self.nchannels)
        self.wav_file.setsampwidth(self.samplewidth)
        self.wav_file.setframerate(self.samplerate)

    def write_frames(self, data: bytes):
        """
        向文件中写入波形
        :param data: 待写入的字节型数据
        :return:
        """
        self.wav_file.writeframes(data)

    def write_float(self, data: float):
        """
        写入一个float类型的波形值
        :param data:
        :return:
        """
        assert self.samplewidth == 2  # TODO: 暂时只支持精度2字节的采样
        self.write_frames(float2sample(data))

    def write_floats(self, data):
        """
        写入一堆float类型的值
        :param data: 可以是ndarray或list类型
        :return:
        """
        datas = bytes()
        for d in data:
            datas += float2sample(d)
        self.write_frames(datas)

    def close(self):
        self.wav_file.close()


class ReadWave:
    def __init__(self, filepath: str):
        """
        :param filepath:
        """
        try:
            with wave.open(filepath, "rb") as wav:
                params = wav.getparams()
                nchannels, sampwidth, framerate, nframes, comptype = params[:5]
                wav_data = wav.readframes(nframes)  # 读取所有音频并返回
        except FileNotFoundError:
            print(f"不存文件'{filepath}'")
            return
        assert nchannels * sampwidth * nframes == len(wav_data)
        self.raw_data = wav_data  # 原始信息
        wav_data = np.frombuffer(wav_data, dtype=get_dtype_by_samplewidth(sampwidth))
        assert len(wav_data) % 2 == 0
        wav_data = wav_data.reshape(-1, nchannels)
        assert wav_data.shape[0] == nframes
        wav_data = wav_data.T  # 每行分别对应一个通道
        self.time_slice = np.arange(0, nframes) / framerate

        self.nchannels = nchannels
        self.sampwidth = sampwidth
        self.framerate = framerate
        self.nframes = nframes
        self.comptype = comptype
        self.wav_data = wav_data
        self.cursor = 0  # 从内存中的raw_data中读取数据的当前值

    def read_frames(self, frame_count):
        """
        读取一帧, 需要注意把帧数转化为字节数
        :param frame_count: 最大读取的帧数
        :return:
        """
        byte_count = frame_count * self.nchannels * self.sampwidth
        if self.cursor + byte_count <= len(self.raw_data):
            next_cursor = self.cursor + byte_count
        else:
            next_cursor = len(self.raw_data)
        ret = self.raw_data[self.cursor: next_cursor]
        self.cursor = next_cursor

        return ret

    def __raw2list(self, chunk: int):
        """
        将bytes类型的raw_data转化为列表, 每个元素是个chunk字节的bytes
        :param chunk: 列表每个元素的字节数
        :return:
        """
        datas = []
        prev_index = 0
        next_index = chunk
        while next_index < len(self.raw_data):
            datas.append(self.raw_data[prev_index:next_index])
            prev_index = next_index
            next_index += chunk
        next_index = len(self.raw_data)
        datas.append(self.raw_data[prev_index:next_index])
        return datas

    def draw_wave(self, internal=1):
        """
        初始化之后直接调用即可绘制频谱图
        :param internal: 用于控制绘制波形的稀疏程度, 采样率太高导致波形几乎是一整个色块
        :return:
        """
        x = self.time_slice[::internal]
        y = self.wav_data[::, ::internal]
        plt.figure()
        if self.nchannels == 1:
            plt.plot(x, y[0])
        elif self.nchannels == 2:
            plt.subplot(211)
            plt.plot(x, y[0])
            plt.subplot(212)
            plt.plot(x, y[1], c="g")
        plt.xlabel("Time/s")
        plt.ylabel("Level")
        plt.show()

    def show_param(self):
        """展示该文件的参数"""
        print(
            f'声道数:{self.nchannels} '
            f'采样字节长度:{self.sampwidth} '
            f'采样频率:{self.framerate} '
            f'总帧数:{self.nframes} '
            f'压缩类型:{self.comptype}'
        )


def main():
    ReadWave("./123")


if __name__ == "__main__":
    main()
